#include<iostream>
using namespace std;
#include<fstream>
#include<cmath>
#include<iomanip>

#include"library.h"

void test(int n, double *d, double *e, double butterfly, int p);

/*
 * tqls: diagonalize a tridiagonal matrix in QL algorithm with shifting
 *
 * int n:		dimension
 * double *d		diagonal elements of the tridiagonal matrix, end up to be the eigen values
 * double *e		subdiagonal elements of the tridiagonal matrix
 * double **v		input as [n][n] matrix, whatever value it has, end up to be the eigen kets
 *
 * convention: e[0]=0
 *
 * d[0]	e[1]
 * e[1]	d[1]
 * ...
 * 			e[n-1]
 * 	  	e[n-1]	d[n-1]
 */
void tqls(int n, double *d, double *e, double **v){

	int m,l,i,iter,j;
	double c,s,r,theta,ks,app,aqq,apq,butterfly;

	// set v as elementary matrix
	for(i=0;i<n;i++){
		for(j=0;j<n;j++){
			if(i==j) v[i][j]=1;
			else v[i][j]=0;
		}
	}

	for(l=0;l<n-1;l++){
		iter=0;//迭代使得e[l+1]足够小，需要的步数
		do{
			for(m=l;m<n-1;m++)//e[m+1]是否已经足够小
				if( abs(e[m+1]) <= 1E-9*(abs(d[m])+abs(d[m+1]))) 
					break;
			//执行完上面的代码块，若e[l+1]足够小，则m=l；若e[l+p]足够小，则m=l+p-1;若都没有足够小，则m=n-1
			//下面的代码处理 l, ..., m 子矩阵块
			if(m!=l){//e[l+1]还没有足够小
				if(iter++ == 30){
					cout<<"Too many iterations in tqli, when trying to decrease e["<<l<<"] l=0,...,n-1 \n";
				}
				theta = (d[l+1]-d[l])/(2.0*e[l+1]);
				ks = - e[l+1]/(fabs(theta) + sqrt(theta*theta+1));
				if( theta < 0 ) ks *= -1;
				ks += d[l];// ks is set up, shifting
				
				for(i=l;i<=m;i++) d[i] -= ks;

				app = d[m-1]; aqq = d[m]; apq = e[m];
				r = sqrt( apq * apq + aqq * aqq);// c * apq - s *aqq = 0
				c = aqq / r;
				s = apq / r;
				// these codes heavily rely on formulas about Jacobi transformation on my Lecture notes, or on Numerical Recipes 
				d[m-1] = c*c* app - 2*c*s*apq + s*s*aqq;
				d[m] = s*s*app + 2*c*s*apq + c*c*aqq;
				e[m] = (c*c-s*s)*apq + c*s*(app-aqq);
				butterfly = s* e[m-1];// a'_{m-2,m} = a'_{m,m-2} = -s a_{m,m-1}
				e[m-1] = c* e[m-1];// a'_{m-2,m-1} = c*a_{m-2,m-1}

				for(j=0;j<n;j++){
					app = v[j][m-1];
					aqq = v[j][m];
					v[j][m-1] = c*app - s*aqq;
					v[j][m] = s*app + c*aqq;
				}

				for(i=m-1;i>=l+1;i--){
					r = sqrt( butterfly*butterfly + e[i+1]*e[i+1]);
					s = butterfly / r;
					c = e[i+1] / r;
					//butterfly = a_{i+1,i-1}, 
					//butterfly' = a'_{i+1,i-1} = c butterfly - s a_{i+1,i} = c butterfly - s e[i+1] = 0

					app = d[i-1]; aqq=d[i]; apq = e[i];
					d[i-1] = c*c* app - 2*c*s*apq + s*s*aqq;
					d[i] = s*s*app + 2*c*s*apq + c*c*aqq;
					e[i] = (c*c-s*s)*apq + c*s*(app-aqq);
					e[i+1] = s*butterfly + c*e[i+1];
					butterfly = s*e[i-1];
					e[i-1] = c*e[i-1];

					for(j=0;j<n;j++){
						app = v[j][i-1];
						aqq = v[j][i];
						v[j][i-1] = c*app - s*aqq;
						v[j][i] = s*app + c*aqq;
					}

				}
				for(i=l;i<=m;i++) d[i] += ks;
			}
		}while( m!=l );
	}
}

void test(int n, double *d, double *e, double butterfly, int p){
	int i,j;
	double *a = new double [n*n];
	double *z = new double [n];

	for(i=0;i<n;i++){
		for(j=i;j<n;j++){
			if(i==j) a[i*n+j] = d[i];
			else if(j==i+1){
				a[i*n+j] = e[j];
				a[j*n+i] = e[j];
			}
			else if(j==p && j==i+2){
				a[i*n+j] = butterfly;
				a[j*n+i] = butterfly;
			}
			else{
			       a[i*n+j] = 0;
			       a[j*n+i] = 0;
			}
		}
	}
	cout<<"a:"<<endl;
	for(i=0;i<n;i++){
		for(j=0;j<n;j++){
			cout<<setw(5)<<setprecision(5)<<a[i*n+j]<<"\t";
		}
		cout<<endl;
	}
	
	jacobi(n,a,z);
	cout<<"eigenvalues: "; for(i=0;i<n;i++)cout<<z[i]<<","; cout<<endl;
	delete [] a,z;
}

int main(){

	int i,j,k,n=10;
	double *d = new double [n];
	double *e = new double [n];
	for(i=0;i<n;i++){
	       d[i] = 10*(i+1);
	       e[i] = i;
	}
	double **v = new double *[n];
	for(i=0;i<n;i++)v[i] = new double [n];
	e[0]=0;

	tqls(n,d,e,v);

        double *a = new double [n*n];
        double *z = new double [n];

        for(i=0;i<n;i++){
                for(j=i;j<n;j++){
                        if(i==j) a[i*n+j] = d[i];
                        else if(j==i+1){
                                a[i*n+j] = e[j];
                                a[j*n+i] = e[j];
                        }
                        else{
                               a[i*n+j] = 0;
                               a[j*n+i] = 0;
                        }
                }
        }

	cout<<"Av"<<endl;
	for(i=0;i<n;i++){
		for(j=0;j<n;j++){
			double y=0;
			for(k=0;k<n;k++){
				y += a[i*n+k] * v[k][j];
			}
			cout<<setw(10)<<setprecision(2)<<y;
		}
		cout<<endl;
	}
	cout<<"v Lambda"<<endl;
	for(i=0;i<n;i++){
		for(j=0;j<n;j++){
			cout<<setw(5)<<setprecision(2)<<d[i]*v[i][j]<<"\t";
		}
		cout<<endl;
	}



	delete [] d,e;

	return 0;
}
